IMAGE=$1

install_pip(){
    pip install --no-cache-dir $@
}

install_pip --upgrade pip
install_pip -r /tmp/requirements.txt

if [[ "$IMAGE" =~ "gpu"  ]];then
    install_pip torch torchvision torchaudio
else
    install_pip torch==1.10.0+cpu torchvision==0.11.1+cpu torchaudio==0.10.0+cpu -f https://download.pytorch.org/whl/cpu/torch_stable.html
fi

